import torch
import argparse

def read_checkpoint(ckpt_path, output_path=None):
    """
    Reads a PyTorch checkpoint file and prints its contents.
    Optionally saves the model state dictionary to a new file.

    Args:
        ckpt_path (str): Path to the checkpoint file.
        output_path (str, optional): Path to save the model state dictionary. Defaults to None.
    """
    # Load the checkpoint file
    try:
        checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
    except FileNotFoundError:
        print(f"Error: The file {ckpt_path} does not exist.")
        return
    except Exception as e:
        print(f"Error loading checkpoint: {e}")
        return

    # Print the contents of the checkpoint
    print("Checkpoint Contents:")
    for key, value in checkpoint.items():
        print(f"{key}: {value}")

    # Optionally save the model state dictionary to a new file
    if output_path:
        try:
            torch.save(checkpoint['state_dict'], output_path)
            print(f"Model state dictionary saved to {output_path}")
        except KeyError:
            print("Error: The checkpoint does not contain a 'state_dict' key.")
        except Exception as e:
            print(f"Error saving model state dictionary: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Read and optionally save a PyTorch checkpoint file.")
    parser.add_argument('--ckpt_path', type=str, required=True, help='Path to the checkpoint file.')
    parser.add_argument('--output_path', type=str, help='Path to save the model state dictionary. Optional.')
    args = parser.parse_args()

    read_checkpoint(args.ckpt_path, args.output_path)